#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>

#define ROOT 0

/*
 * Aloca uma matriz na forma um vetor de tamanho rows * cols e retorna um ponteiro para um vetor
 * 	de forma que este possa ser acessado como as posicoes de uma matriz.
 */
float **alloc_2d_float(int rows, int cols) {
	int i;

	float *data = (float *) malloc(rows * cols * sizeof(float));
	float **array = (float **) malloc(rows * sizeof(float*));
	for (i = 0; i < rows; i++)
		array[i] = &(data[cols * i]);

	return array;
}

/*
 * Desaloca a memoria dos vetores alocados pela funcao acima.
 */
void desalloc_2d_float(float **matrix) {
	free(matrix[0]);
	free(matrix);
}

int main(int argc, char *argv[]) {
	float **matriz;
	float **matrizAEstrela;
	float **localMatrix;
	float *matrizB;
	float *matrizBEstrela;
	float *X0;
	float *auxMultiplicacao;
	float *X1;
	float *result, *global_result;
	int i, j, k, iteracao = 1, converged = 0;
	float row_test = 0.0;
	int *send_counts, *recv_cnt, *displs, *recv_displs, chunks_left;

	int rank, size, chunk_size;
	double startTime;

	int J_ORDER;
	int J_ROW_TEST;
	int J_ITE_MAX;
	float J_ERROR;

	MPI_Init( &argc,&argv); /* Inicia o ambiente paralelo */
	MPI_Comm_rank( MPI_COMM_WORLD, &rank); /* Retorna o rank do processo */
	MPI_Comm_size( MPI_COMM_WORLD, &size); /* Retorna o número de processos no comunicador */

	/* Abre arquivo para leitura dos dados do sistema a ser resolvido */
	FILE *fp = fopen(argv[1], "r");
	if (fp == NULL) {
		printf("Error: loading file!\n");
		exit(EXIT_FAILURE);
	}
	rewind(fp);
	/* Le ordem da matriz, linha para teste de convergencia, erro e
		num maximo de iteracoes, respectivamente */
	fscanf(fp, "%d%d%f%d", &J_ORDER, &J_ROW_TEST, &J_ERROR, &J_ITE_MAX);

	/* Alocação das matrizes e vetores */
	matrizB = (float *) malloc(J_ORDER * sizeof(float));
	matrizBEstrela = (float *) malloc(J_ORDER * sizeof(float));
	X0 = (float *) malloc(J_ORDER * sizeof(float));
	auxMultiplicacao = (float *) malloc(J_ORDER * sizeof(float));
	X1 = (float *) malloc(J_ORDER * sizeof(float));
	result = (float *) malloc(J_ORDER * sizeof(float));
	global_result = (float *) malloc(J_ORDER * sizeof(float));
	send_counts = (int *) malloc(size * sizeof(int));
	displs = (int *) malloc(size * sizeof(int));
	recv_cnt = (int *) malloc(size * sizeof(int));
	recv_displs = (int *) malloc(size * sizeof(int));

	matriz = alloc_2d_float(J_ORDER, J_ORDER);
	matrizAEstrela = alloc_2d_float(J_ORDER, J_ORDER);
	localMatrix = alloc_2d_float(J_ORDER, J_ORDER);

	chunk_size = J_ORDER / size; /* Define o tamanho dos blocos que serao enviados para os processos */

	chunks_left = J_ORDER; /* Decrementado para o caso de J_ORDER nao ser uma potencia de 'size' (ultimo bloco enviado tera tamanho menor) */

	for (i = 0; i < size; i++) {
		recv_cnt[i] = chunk_size; /* Elementos a serem reunidos de cada processo */
		recv_displs[i] = i * chunk_size; /* Posicao no recev_buffer dos elementos reunidos de cada processo */
	}

	/* Procedimentos realizados apenas pelo nó mestre */
	if (rank == 0) {
		
		startTime = MPI_Wtime();

		printf("Reading matrices... ");

		/* Leitura da matriz de coeficientes */
		for (i = 0; i < J_ORDER; i++)
			for (j = 0; j < J_ORDER; j++) {
				fscanf(fp, "%f", &matriz[i][j]);
				matrizAEstrela[i][j] = matriz[i][j];
			}

		/* Leitura do vetor B */
		for (i = 0; i < J_ORDER; i++) {
			fscanf(fp, "%f", &matrizB[i]);
			matrizBEstrela[i] = matrizB[i];
		}

		fclose(fp); /* Fecha arquivo */

		printf("Done.\nCalculating... ");

		/*
		 * Inicio da paralelizacao com OpenMP.
		 */
		#pragma omp parallel shared (matriz, matrizAEstrela, matrizBEstrela) private (i,j)
		{
			#pragma omp for schedule (static)
			/* Divisao do vetor B pela respectiva diagonal da matriz A */
			for (i = 0; i < J_ORDER; i++) {
				matrizBEstrela[i] /= matriz[i][i];
				X0[i] = matrizBEstrela[i];
			}

			/* Calculo da matriz A* */
			#pragma omp for schedule (static)
			for (i = 0; i < J_ORDER; i++) {
				for (j = 0; j < J_ORDER; j++) {
					if (i != j) {
						matrizAEstrela[i][j] /= matriz[i][i];
						matrizAEstrela[i][j] *= -1;
					}
					matrizAEstrela[i][i] = 0;
				}
			}
		}

		/* Define o tamanho dos blocos que serao enviados para cada processo */
		for (i = 0; i < size; i++) {
			/* Se a ordem da matriz nao for um multiplo do numero de nós o ultimo processo recebera um bloco menor */
			if ((chunks_left - chunk_size) > 0)
	                	send_counts[i] = chunk_size * J_ORDER;
			else
		                send_counts[i] = chunks_left * J_ORDER; /* O bloco menor eh definido aqui */
			chunks_left -= chunk_size;
			displs[i] = i * send_counts[i]; /* Semelhante ao recv_displs porem para a funcao MPI_Scatterv */
		}

		/* Nó mestre faz broadcast de X0 */
		MPI_Bcast(X0, J_ORDER, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

		/* Nó mestre faz broadcast do vetor com o tamanho dos blocos enviados para cada processo */
		MPI_Bcast(send_counts, size, MPI_INT, ROOT, MPI_COMM_WORLD);

		/* Nó mestre envia partes da matriz A* para cada nó */
		MPI_Scatterv(&(matrizAEstrela[0][0]), send_counts, displs, MPI_FLOAT, &(localMatrix[0][0]), chunk_size * J_ORDER, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

		for (k = 0; k < J_ITE_MAX && !converged; k++) {

			/* Multiplicacao da sua parte de A* (linhas) pelo vetor B* */
			for (i = 0; i < send_counts[rank] / J_ORDER; i++) {
				result[i] = 0;
				for (j = 0; j < J_ORDER; j++)
					result[i] += localMatrix[i][j] * X0[j];
			}

			/* Processo mestre faz Gather do resultado */
			MPI_Gatherv(result, chunk_size, MPI_FLOAT, global_result, recv_cnt, recv_displs, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

			/* Calcula X[k+1] */
			for (i = 0; i < J_ORDER; i++)
				X1[i] = global_result[i] + matrizBEstrela[i];

			/* Teste de convergencia (J_ROW_TEST) */
			row_test = 0;
			for (i = 0; i < J_ORDER; i++)
				row_test += matriz[J_ROW_TEST][i] * X1[i];

			/* Atualiza X0 (X[k]) para a proxima iteracao */
			for (i = 0; i < J_ORDER; i++)
				X0[i] = X1[i];

			iteracao++;

			/* Faz broadcast do vetor atualizado */
			MPI_Bcast(X0, J_ORDER, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

			/* Verifica convergencia de acordo com o erro */
			if (fabs(matrizB[J_ROW_TEST] - row_test) < J_ERROR)
				converged = 1;

			/* Faz broadcast do flag de convergencia (criterio de parada) */
			MPI_Bcast(&converged, 1, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

		}

		printf("Done.\nEXECUTION TIME: %lf seconds\n\n", MPI_Wtime() - startTime);

		printf("---------------------------------------------------------\n");
		printf("Iterations: %d\n", iteracao);
		printf("RowTest: %d => [%f] =? %f\n", J_ROW_TEST, row_test, matrizB[J_ROW_TEST]);
		printf("---------------------------------------------------------\n");

	}
	else { /* Nó escravo */

		/* Nós escravos recebem o broadcast do vetor X0 */
		MPI_Bcast(X0, J_ORDER, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

		/* Nós escravos recebem o broadcast do vetor send_counts */
		MPI_Bcast(send_counts, size, MPI_INT, ROOT, MPI_COMM_WORLD);

		/* Nós escravos recebem as partes da matriz A* */
		MPI_Scatterv(&(matrizAEstrela[0][0]), send_counts, displs, MPI_FLOAT, &(localMatrix[0][0]), chunk_size * J_ORDER, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

		for (k = 0; k < J_ITE_MAX && !converged; k++) {

			/* Multiplicacao da sua parte de A* (linhas) pelo vetor B* */
			for (i = 0; i < send_counts[rank] / J_ORDER; i++) {
				result[i] = 0;
				for (j = 0; j < J_ORDER; j++)
					result[i] += localMatrix[i][j] * X0[j];
			}

			/* Processos escravos fazem Gather do resultado */
			MPI_Gatherv(result, chunk_size, MPI_FLOAT, global_result, recv_cnt, recv_displs, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

			/* Recebe broadcast do vetor atualizado */
			MPI_Bcast(X0, J_ORDER, MPI_FLOAT, ROOT, MPI_COMM_WORLD);

			/* Recebe broadcast do flag de convergencia (criterio de parada) */
			MPI_Bcast(&converged, 1, MPI_FLOAT, ROOT, MPI_COMM_WORLD);
		}

	}

	/* Libera espacos alocados para matrizes e vetores */
	desalloc_2d_float(matriz);
	desalloc_2d_float(matrizAEstrela);
	desalloc_2d_float(localMatrix);
	free(matrizB);
	free(matrizBEstrela);
	free(X0);
	free(auxMultiplicacao);
	free(X1);
	free(result);
	free(global_result);
	free(send_counts);
	free(recv_cnt);
	free(displs);
	free(recv_displs);

	MPI_Finalize(); /* Finaliza o ambiente de execucao da MPI. Todo processo deve chamar esta funcao antes de finalizar */
	return 0;
}

